from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, OperationFailure
from tqdm import tqdm
import time
import logging

# 假设您有一个 log_config.py，如果没有，可以使用基础配置
# from config.log_config import get_logger
# logger = get_logger()

# --- 基础日志配置 (如果没有 log_config.py) ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# --- 配置区 ---
# MongoDB 连接信息
MONGO_URI = "mongodb://localhost:27017/"
DATABASE_NAME = "LocalArkTSTest"

# 集合名称
LIBRARY_INFO_COLLECTION = "library_info_testv2"
NPM_FUNCTIONS_COLLECTION = "npm_functions_testv2"


def update_library_info_batch():
    """
    高效地更新 library_info 集合中的 versionInfo 数据。
    此版本兼容 version 字段为字符串或字符串数组的情况。
    此函数会为 versionInfo 数组中的每个对象填充 'func_num' 和 'files' 字段。
    """
    start_time = time.time()
    logger.info("开始第一步：更新 versionInfo 中的 func_num 和 files...")

    try:
        # 1. 连接到 MongoDB
        client = MongoClient(MONGO_URI)
        client.admin.command('ping')  # 验证连接
        logger.info("MongoDB 连接成功！")

        db = client[DATABASE_NAME]
        library_info_col = db[LIBRARY_INFO_COLLECTION]
        npm_functions_col = db[NPM_FUNCTIONS_COLLECTION]

        # 2. 获取所有需要处理的 library_info 文档
        libraries_to_process = list(library_info_col.find({}, {"_id": 1, "library": 1}))
        total_libraries = len(libraries_to_process)

        if total_libraries == 0:
            logger.info(f"在 '{LIBRARY_INFO_COLLECTION}' 集合中未找到任何文档。")
            return

        logger.info(f"找到 {total_libraries} 个库需要处理。")

        # 3. 遍历每个库并执行更新
        for lib_info_doc in tqdm(libraries_to_process, desc="Processing Libraries"):
            lib_name = lib_info_doc.get('library')
            if not lib_name:
                logger.warning(f"警告: 跳过没有 'library' 字段的文档 (ID: {lib_info_doc['_id']})")
                continue

            # 步骤 A: 一次性从 npm_functions 集合中获取该库的所有函数
            all_functions = list(npm_functions_col.find({"library": lib_name}))

            # 步骤 B: 在内存中预处理数据，按版本分组统计
            version_stats = {}
            for func in all_functions:
                version_value = func.get('version')
                relative_path = func.get('relativePath')

                if not relative_path:
                    continue

                versions_to_process = []
                if isinstance(version_value, list):
                    versions_to_process = version_value
                elif isinstance(version_value, str):
                    versions_to_process = [version_value]
                else:
                    continue

                if not versions_to_process:
                    continue

                for version in versions_to_process:
                    if version not in version_stats:
                        version_stats[version] = {'func_num': 0, 'files': set()}
                    version_stats[version]['func_num'] += 1
                    version_stats[version]['files'].add(relative_path)

            # 步骤 C: 获取完整的 library_info 文档以进行修改
            full_lib_info_doc = library_info_col.find_one({"_id": lib_info_doc["_id"]})
            if not full_lib_info_doc or 'versionInfo' not in full_lib_info_doc:
                logger.warning(f"警告: 无法找到库 '{lib_name}' 的完整文档或其缺少 'versionInfo' 字段。")
                continue

            # 步骤 D: 遍历 versionInfo 数组，使用预计算的数据进行更新
            updated_version_info = []
            for version_obj in full_lib_info_doc.get('versionInfo', []):
                version = version_obj.get('version')
                if not version:
                    continue

                stats = version_stats.get(version)
                if stats:
                    version_obj['func_num'] = stats['func_num']
                    version_obj['files'] = len(stats['files'])
                else:
                    version_obj['func_num'] = 0
                    version_obj['files'] = 0

                updated_version_info.append(version_obj)

            # 步骤 E: 执行单次更新操作，将整个 versionInfo 数组写回
            try:
                library_info_col.update_one(
                    {"_id": lib_info_doc["_id"]},
                    {"$set": {"versionInfo": updated_version_info}}
                )
            except OperationFailure as e:
                logger.error(f"错误: 更新库 '{lib_name}' 的 versionInfo 失败: {e}")

    except ConnectionFailure:
        logger.error("错误: 无法连接到 MongoDB。请检查服务是否运行以及连接字符串是否正确。")
    except Exception as e:
        logger.error(f"发生未知错误: {e}", exc_info=True)
    finally:
        if 'client' in locals() and client:
            client.close()
            end_time = time.time()
            logger.info(f"第一步处理完成！总耗时: {end_time - start_time:.2f} 秒。")
            logger.info("MongoDB 连接已关闭。")


def calculate_and_update_summary_stats():
    """
    第二步：根据已填充的 versionInfo 数据，计算并更新每个库的
    total_function_number 和 average_number 字段。
    """
    start_time = time.time()
    logger.info("开始第二步：计算并更新 total_function_number 和 average_number...")

    try:
        client = MongoClient(MONGO_URI)
        client.admin.command('ping')
        logger.info("MongoDB 连接成功！")

        db = client[DATABASE_NAME]
        library_info_col = db[LIBRARY_INFO_COLLECTION]

        # 获取所有库文档
        libraries_to_process = list(library_info_col.find({}))
        total_libraries = len(libraries_to_process)

        if total_libraries == 0:
            logger.info(f"在 '{LIBRARY_INFO_COLLECTION}' 集合中未找到任何文档。")
            return

        logger.info(f"找到 {total_libraries} 个库需要计算摘要信息。")

        for lib_doc in tqdm(libraries_to_process, desc="Calculating Summary Stats"):
            version_info = lib_doc.get('versionInfo', [])

            if not version_info:
                # 如果 versionInfo 为空或不存在，则将统计值设为 0
                total_function_number = 0
                average_number = 0
            else:
                # 计算所有版本 func_num 的总和
                total_function_number = sum(item.get('func_num', 0) for item in version_info)

                # 计算平均值
                version_count = len(version_info)
                average_number = total_function_number / version_count if version_count > 0 else 0

            # 执行更新操作
            try:
                library_info_col.update_one(
                    {"_id": lib_doc["_id"]},
                    {
                        "$set": {
                            "total_function_number": total_function_number,
                            "average_number": average_number
                        }
                    }
                )
            except OperationFailure as e:
                logger.error(f"错误: 更新库 '{lib_doc.get('library', 'Unknown')}' 的摘要信息失败: {e}")

    except ConnectionFailure:
        logger.error("错误: 无法连接到 MongoDB。")
    except Exception as e:
        logger.error(f"发生未知错误: {e}", exc_info=True)
    finally:
        if 'client' in locals() and client:
            client.close()
            end_time = time.time()
            logger.info(f"第二步处理完成！总耗时: {end_time - start_time:.2f} 秒。")
            logger.info("MongoDB 连接已关闭。")


if __name__ == "__main__":
    # 按顺序执行两个步骤
    # 第一步：填充 versionInfo 中的 func_num 和 files
    update_library_info_batch()
    # 第二步：基于第一步的结果，计算并更新 total_function_number 和 average_number
    calculate_and_update_summary_stats()
